﻿//==============================================================================
// File: ManagedThreadPool.cs
// Created: 2010-06-07
// Author: Piotr Włodek, Arkadiusz Świerczek
//==============================================================================
// This file is a part of MAMMOTH project.
// Copyright (C) 2010 AGH University of Science and Technology, Krakow.
// https://caribou.iisg.agh.edu.pl/trac/mammoth
//
// Mammoth is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
// 
// Mammoth is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Lesser General Public License for more details.
// 
// You should have received a copy of the GNU Lesser General Public License
// along with Mammoth. If not, see http://www.gnu.org/licenses/.
//==============================================================================

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Threading;
using WpfRcp.Core.Common;
using ThreadState = System.Threading.ThreadState;

namespace WpfRcp.Core.Threading
{
    /// <summary>
    /// Represents simple, managed pool of threads.
    /// </summary>
    public sealed class ManagedThreadPool
    {
        /// <summary>
        /// Event raised when a task is finished or cancelled.
        /// </summary>
        public event EventHandler<TaskCompletedEventArgs> TaskCompleted;

        /// <summary>
        /// Event raised when an exception occures in a task.
        /// </summary>
        public event EventHandler<TaskErrorEventArgs> TaskError;

        /// <summary>
        /// Event raised whenever worker thread becomes idle.
        /// </summary>
        public event EventHandler<EventArgs> ThreadIdle;

        /// <summary>
        /// Default pool size.
        /// </summary>
        public static readonly int DefaultPoolSize = Environment.ProcessorCount;

        // Fields
        private readonly Queue<WorkerThread> m_WorkerThreads;
        private readonly Queue<TaskResult> m_Tasks;

        /// <summary>
        /// Constructs pool with its default size.
        /// </summary>
        public ManagedThreadPool() : this(DefaultPoolSize) { }

        /// <summary>
        /// Constructs pool with specified number of threads.
        /// </summary>
        /// <param name="poolSize">Number of threads in the pool.</param>
        public ManagedThreadPool(int poolSize)
        {
            m_WorkerThreads = new Queue<WorkerThread>();
            m_Tasks = new Queue<TaskResult>();
            WorkersCount = poolSize;

            // lock statement can be omitted here
            for (int i = 0; i < WorkersCount; i++)
            {
                m_WorkerThreads.Enqueue(new WorkerThread(this));
            }
        }

        /// <summary>
        /// Gets the number of workers in the pool.
        /// </summary>
        public int WorkersCount { get; private set; }

        /// <summary>
        /// Gets the value indicating whether the pool has been stopped.
        /// </summary>
        public bool Stopped { get; private set; }

        /// <summary>
        /// Gets a count of unprocessed tasks.
        /// </summary>
        public int TaskCount
        {
            get
            {
                lock (m_Tasks)
                {
                    return m_Tasks.Count;
                }
            }
        }

        /// <summary>
        /// Gets a count of idle workers.
        /// </summary>
        public int IdleWorkersCount
        {
            get
            {
                lock (m_WorkerThreads)
                {
                    return m_WorkerThreads.Count;
                }
            }
        }

        /// <summary>
        /// Adds new workers to the pool.
        /// </summary>
        /// <param name="num">Number of workers to be added.</param>
        public void AddWorkers(int num)
        {
            ThrowIfStopped();

            lock (m_WorkerThreads)
            {
                WorkersCount += num;
                for (int i = 0; i < num; i++)
                {
                    m_WorkerThreads.Enqueue(new WorkerThread(this));
                }
            }
        }

        /// <summary>
        /// Stops the thread pool. Once stopped, the pool cannot be reused.
        /// </summary>
        public void Stop()
        {
            foreach (var worker in m_WorkerThreads)
            {
                lock (worker)
                {
                    Monitor.Pulse(worker);
                }
            }
            Stopped = true;
        }

        /// <summary>
        /// Schedules user task to be executed in the thread bool when any worker is available.
        /// </summary>
        /// <param name="task">Task to be executed.</param>
        /// <returns>Returns <see cref="ITaskResult"/> object that contains
        /// the state of the task's execution.</returns>
        public ITaskResult QueueUserWorkItem(ITask task)
        {
            ThrowIfStopped();

            TaskResult result;

            lock (m_Tasks)
            lock (m_WorkerThreads)
            {
                result = new TaskResult { Task = task };

                if (m_WorkerThreads.Count > 0)
                {
                    // There are free workers, schedule the execution
                    m_WorkerThreads.Dequeue().SetTask(result);
                }
                else
                {
                    // No free workers, save the task
                    m_Tasks.Enqueue(result);
                }
            }

            return result;
        }

        public ITaskResult QueueUserWorkItem(Action action)
        {
            return QueueUserWorkItem(new RunnableTask(action));
        }

        public ITaskResult QueueUserWorkItem<T>(Action<T> action) where T : class
        {
            return QueueUserWorkItem(action, null);
        }

        public ITaskResult QueueUserWorkItem<T>(Action<T> action, T param)
        {
            return QueueUserWorkItem(new RunnableTask<T>(action, param));
        }

        private void ReturnWorker(WorkerThread worker)
        {
            lock (m_Tasks)
            lock (m_WorkerThreads)
            {
                if (m_Tasks.Count > 0)
                {
                    worker.SetTask(m_Tasks.Dequeue());
                }
                else
                {
                    m_WorkerThreads.Enqueue(worker);
                    OnThreadIdle(this, EventArgs.Empty);
                }
            }
        }

        private void OnWorkCompleted(object sender, TaskCompletedEventArgs e)
        {
            EventsHelper.Fire(TaskCompleted, sender, e);
        }

        private void OnTaskError(object sender, TaskErrorEventArgs e)
        {
            EventsHelper.Fire(TaskError, sender, e);
        }

        private void OnThreadIdle(object sender, EventArgs e)
        {
            EventsHelper.Fire(ThreadIdle, sender, e);
        }

        private void ThrowIfStopped()
        {
            if (Stopped)
                throw new InvalidOperationException("Thread pool is stopped.");
        }

        private sealed class WorkerThread
        {
            /// <summary>
            /// Worker thread name.
            /// </summary>
            private const string WorkerThreadName = "ThreadPool Worker Thread";

            // Fields
            private readonly Thread m_Thread;
            private readonly ManagedThreadPool m_ThreadPool;
            private TaskResult m_TaskResult;

            public WorkerThread(ManagedThreadPool threadPool)
            {
                m_ThreadPool = threadPool;
                m_Thread = new Thread(Run) { Name = WorkerThreadName, IsBackground = true };
                m_Thread.Start();
            }

            public Thread Thread
            {
                get { return m_Thread; }
            }

            /// <summary>
            /// Gets a value indicating if the thread pool is stopped.
            /// </summary>
            public bool Stopped { get; private set; }

            public void SetTask(TaskResult task)
            {
                m_TaskResult = task;
                m_TaskResult.CurrentWorker = this;

                if (task.Status == TaskStatus.Cancelled)
                {
                    m_TaskResult = null;
                    m_ThreadPool.ReturnWorker(this);

                    return;
                }

                lock (this)
                {
                    Monitor.Pulse(this);
                }
            }

            private void Run()
            {
                while (m_ThreadPool.Stopped == false)
                {
                    lock (this)
                    {
                        while (m_TaskResult != null)
                        {
                            try
                            {
                                m_TaskResult.Status = TaskStatus.Running;
                                var returnVal = m_TaskResult.Task.Run();

                                // Run has completed so we cannot abort thread before we
                                // raise the event and set the status
                                Thread.BeginCriticalRegion();

                                m_TaskResult.Status = TaskStatus.Completed;
                                m_ThreadPool.OnWorkCompleted(m_ThreadPool, new TaskCompletedEventArgs(m_TaskResult.Task, returnVal));
                                m_TaskResult = null;

                                Thread.EndCriticalRegion();
                            }
                            catch (ThreadAbortException)
                            {
                                Debug.Assert(m_TaskResult != null);

                                // If status is completed and we still got an error,
                                // it seems OnWorkCompleted thrown a ThreadAbortException so 
                                // Thread.EndCriticalRegion() didn't have a chance to be executed, do it now
                                if (m_TaskResult.Status == TaskStatus.Completed)
                                    Thread.EndCriticalRegion();

                                m_TaskResult.Status = TaskStatus.Cancelled;
                                m_ThreadPool.OnWorkCompleted(m_ThreadPool, new TaskCompletedEventArgs(m_TaskResult.Task, null));
                                m_TaskResult = null;

                                // Restore worker to its normal state
                                Thread.ResetAbort();
                            }
                            catch (Exception ex)
                            {
                                Debug.Assert(m_TaskResult != null);

                                // If status is completed and we still got an error,
                                // it seems OnWorkCompleted thrown an exception so 
                                // Thread.EndCriticalRegion() didn't have a chance to be executed, do it now
                                if (m_TaskResult.Status == TaskStatus.Completed)
                                    Thread.EndCriticalRegion();

                                // This is done because it is possible that the ThreadAbortException has been thrown, caught,
                                // and another exception has been thrown
                                if ((Thread.CurrentThread.ThreadState  & ThreadState.AbortRequested) == ThreadState.AbortRequested)
                                {
                                    m_TaskResult.Status = TaskStatus.Cancelled;
                                    m_ThreadPool.OnWorkCompleted(m_ThreadPool, new TaskCompletedEventArgs(m_TaskResult.Task, null));
                                    
                                    // Restore worker to its normal state
                                    Thread.ResetAbort();
                                }
                                else
                                {
                                    m_TaskResult.Exception = ex;
                                    m_TaskResult.Status = TaskStatus.Error;
                                    m_ThreadPool.OnTaskError(this, new TaskErrorEventArgs(m_TaskResult.Task, ex));
                                }

                                m_TaskResult = null;
                            }

                            m_ThreadPool.ReturnWorker(this);
                        }

                        // Wait for another task
                        Monitor.Wait(this);
                    }
                }

                Stopped = true;
            }
        }

        private class RunnableTask<T> : ITask
        {
            private readonly Action<T> m_Action;
            private readonly T m_Value;

            public RunnableTask(Action<T> action, T param)
            {
                m_Action = action;
                m_Value = param;
            }

            public object Run()
            {
                if (m_Action != null)
                {
                    m_Action(m_Value);
                }
                return null;
            }
        }

        private class RunnableTask : ITask
        {
            private readonly Action m_Action;

            public RunnableTask(Action action)
            {
                m_Action = action;
            }

            public object Run()
            {
                if (m_Action != null)
                {
                    m_Action();
                }
                return null;
            }
        }

        private class TaskResult : ITaskResult
        {
            // Private fields
            private TaskStatus m_Status;

            // Events
            public event EventHandler<EventArgs> StateChanged;

            // Public properties
            public ManualResetEvent Handle { get; private set; }
            public ITask Task { get; set; }
            public WorkerThread CurrentWorker { get; set; }
            public Exception Exception { get; set; }

            public TaskStatus Status
            {
                get { return m_Status; }
                set
                {
                    m_Status = value;
                    if (value != TaskStatus.Queued && value != TaskStatus.Running)
                        Handle.Set();

                    if (StateChanged != null)
                        StateChanged(this, EventArgs.Empty);
                }
            }

            public TaskResult()
            {
                Handle = new ManualResetEvent(false);
            }

            public void Cancel()
            {
                // If queued, cancel before the task is scheduled on the CPU
                if (Status == TaskStatus.Queued)
                    Status = TaskStatus.Cancelled;
                else if (Status != TaskStatus.Running)
                    throw new ApplicationException("Cannot cancel finished task.");
                else
                    // Task is running, abort it
                    CurrentWorker.Thread.Abort();
            }

            public bool IsCompleted
            {
                get
                {
                    return Status == TaskStatus.Completed ||
                           Status == TaskStatus.Cancelled || Status == TaskStatus.Error;
                }
            }

            public WaitHandle AsyncWaitHandle
            {
                get { return Handle; }
            }

            public object AsyncState
            {
                get { return null; }
            }

            public bool CompletedSynchronously
            {
                get { return false; }
            }
        }
    }

    /// <summary>
    /// Represents a task result.
    /// </summary>
    public interface ITaskResult : IAsyncResult
    {
        /// <summary>
        /// Event that is raised whenever task state is changed.
        /// </summary>
        event EventHandler<EventArgs> StateChanged;

        /// <summary>
        /// Gets a task status.
        /// </summary>
        TaskStatus Status { get; }

        /// <summary>
        /// Gets an exception that occured during task execution. Returns null if no error occured.
        /// </summary>
        Exception Exception { get; }

        /// <summary>
        /// Cancels the task.
        /// </summary>
        void Cancel();
    }

    /// <summary>
    /// Represents a task status.
    /// </summary>
    public enum TaskStatus
    {
        Queued,
        Running,
        Completed,
        Cancelled,
        Error
    }
}